package com.wan.net.rpc.invoker.callback;

import com.wan.net.exception.RpcInvokeException;

import java.util.concurrent.*;
import java.util.function.Consumer;

/**
 * rpc请求的结果持有
 * @author wwh
 * @date 2022/6/22 17:18
 */
public class RpcFuture<R> implements Future<R> {

    private R response;
    private CountDownLatch lock;
    private boolean cancel = false;
    private final Consumer<R> consumer;
    private RuntimeException exception;

    public RpcFuture() {
        this(null);
    }

    public RpcFuture(Consumer<R> consumer){
        this.lock = new CountDownLatch(1);
        this.consumer = consumer;
    }


    @Override
    public boolean cancel(boolean mayInterruptIfRunning) {
        lock.countDown();
        lock = null;
        response = null;
        cancel = true;
        return true;
    }

    @Override
    public boolean isCancelled() {
        return cancel;
    }

    @Override
    public boolean isDone() {
        return !isCancelled() && response != null && lock.getCount() == 0;
    }

    @Override
    public R get() throws InterruptedException, ExecutionException {
        tryThrowException();
        if(isCancelled() || isDone()){
            return response;
        }
        lock.await();
        tryThrowException();
        return response;
    }

    @Override
    public R get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
        tryThrowException();
        if(isCancelled() || isDone()){
            return response;
        }
        if(lock.await(timeout, unit)){
            tryThrowException();
            return response;
        }
        throw new RpcInvokeException("未知","获取结果超时",new TimeoutException());
    }

    /**
     * 异常
     */
    public void error(RuntimeException exception){
        this.exception = exception;
        lock.countDown();
    }

    private void tryThrowException() {
        if(exception != null){
            throw exception;
        }
    }

    /**
     * 任务完成
     */
    public void complete(R response){
        this.response = response;
        lock.countDown();
        if(consumer != null){
            consumer.accept(response);
        }
    }
}
